package algorithm.encode.huffman;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;

public class HuffmanDictionary {
    String input;
    HuffmanTreeNode root;
    Map<String, String> dictionary;
    Map<String, String> codeDictionary;

    public HuffmanDictionary(String input) {
        this.input = input;
    }


    public void build() {
        List<HuffmanTreeNode> nodeList = buildNodeList(input);
        root = buildHuffmanTreeNode(nodeList);
        prepareCode(root);
        dictionary = createDictionary(root);
        codeDictionary = createCodeDictionary(dictionary);
    }

    private Map<String, String> createCodeDictionary(Map<String, String> dictionary) {
        Map<String, String> codeDictionary = new HashMap<>();
        dictionary.forEach((key, value) -> {
            codeDictionary.put(value, key);
        });
        return codeDictionary;
    }

    private Map<String, String> createDictionary(HuffmanTreeNode root) {
        Map<String, String> dictionary = new HashMap<>();
        visit(root, (HuffmanTreeNode node) -> {
            if (node.data != null) {
                dictionary.put(node.data, node.code);
            }
        });
        return dictionary;
    }

    private void visit(HuffmanTreeNode root, Consumer<HuffmanTreeNode> handler) {
        if (root != null) {
            handler.accept(root);
            visit(root.left, handler);
            visit(root.right, handler);
        }
    }

    private void prepareCode(HuffmanTreeNode root) {
        prepareCode(root, "");
    }

    private void prepareCode(HuffmanTreeNode node, String prefix) {
        if (node != null) {
            node.code = prefix;
            prepareCode(node.left, prefix + "0");
            prepareCode(node.right, prefix + "1");
        }
    }

    private HuffmanTreeNode buildHuffmanTreeNode(List<HuffmanTreeNode> nodeList) {
        if (nodeList.size() == 0) {
            return null;
        } else if (nodeList.size() == 1) {
            HuffmanTreeNode left = nodeList.remove(0);
            HuffmanTreeNode parent = new HuffmanTreeNode();
            parent.weight = left.weight;
            parent.setLeft(left);
            return parent;
        } else {
            while (nodeList.size() > 0) {
                if (nodeList.size() == 1) {
                    break;
                } else {
                    nodeList.sort((o1, o2) -> o1.weight - o2.weight);
                    HuffmanTreeNode left = nodeList.remove(0);
                    HuffmanTreeNode right = nodeList.remove(0);
                    HuffmanTreeNode parent = new HuffmanTreeNode();
                    parent.weight = left.weight + right.weight;
                    parent.setLeft(left);
                    parent.setRight(right);
                    nodeList.add(parent);
                }
            }
            return nodeList.remove(0);
        }
    }

    private List<HuffmanTreeNode> buildNodeList(String input) {
        Map<String, Integer> nodeInfoMap = new HashMap<>();
        String[] inputChars = input.split("");
        for (int i = 0; i < inputChars.length; i++) {
            nodeInfoMap.compute(inputChars[i], (key, value) -> value == null ? 1 : value + 1);
        }
        List<HuffmanTreeNode> nodeList = new ArrayList<>();
        nodeInfoMap.forEach((key, value) -> {
            HuffmanTreeNode node = new HuffmanTreeNode();
            node.weight = value;
            node.data = key;
            nodeList.add(node);
        });
        return nodeList;
    }
    public void printDictionary(){
        dictionary.forEach((key,value)->{
            System.out.println(String.format("%5s -> %5s", key, value));
        });
    }
    public String getInput() {
        return input;
    }

    public HuffmanTreeNode getRoot() {
        return root;
    }

    public Map<String, String> getDictionary() {
        return dictionary;
    }
    public Map<String, String> getCodeDictionary() {
        return codeDictionary;
    }
}
